package edu.washington.cs.sounddetector.server.weka;

import java.util.Collection;
import java.util.List;

import javax.jdo.PersistenceManager;

import weka.classifiers.Classifier;
import weka.core.Attribute;
import weka.core.FastVector;
import weka.core.Instance;
import weka.core.Instances;
import edu.washington.cs.sounddetector.server.database.Model;
import edu.washington.cs.sounddetector.server.database.Sound;
import edu.washington.cs.sounddetector.server.database.DbSoundClass;

/**
 * Static helper methods for dealing with Weka related tasks.
 * 
 * @author the.dylan.price@gmail.com
 * 
 */
public class WekaHelper
{

    public static FastVector createAttributesVector(Collection<String> soundClasses,
            Collection<String> attributes)
    {
        FastVector attInfo = new FastVector(attributes.size());

        // add class attribute
        FastVector fvNominalVal = new FastVector();
        for (String soundClass : soundClasses)
            fvNominalVal.addElement(soundClass);

        Attribute classAttr = new Attribute(WekaConsts.SOUND_CLASS_NAME, fvNominalVal);
        attInfo.addElement(classAttr);

        // add features attributes
        for (String attribute : attributes)
            attInfo.addElement(new Attribute(attribute));

        return attInfo;
    }

    /**
     * Populates an {@link Instances}.
     * 
     * @param sounds
     *            a collection of sounds to create instances for
     * @param attInfo
     *            a vector of attribute information (as returned by
     *            {@link #createAttributesVector(Collection, List)}.
     * @return a new {@link Instances} populated with and {@link Instance} per sound
     */
    public static Instances getInstances(Collection<Sound> sounds, FastVector attInfo)
    {

        // create instances
        Instances instances = new Instances(WekaConsts.RELATION_NAME, attInfo, sounds.size());
        instances.setClassIndex(0);

        for (Sound sound : sounds)
        {
            Instance instance = new Instance(1.0, sound.getFeatures());
            instance.insertAttributeAt(0);
            instance.setDataset(instances);
            instance.setClassValue(sound.getSoundClass().getSoundClass());

            instances.add(instance);
        }

        return instances;
    }

    public static void trainModel(PersistenceManager pm, Model model) throws Exception
    {
        String detectorId = model.getDetectorId();
        // set up instances
        // TODO: change when we implement users
        Collection<DbSoundClass> soundClasses = DbSoundClass.query(pm, "1", detectorId);
        Collection<String> soundClassNames = DbSoundClass.extractSoundClassNames(soundClasses);

        FastVector attInfo = WekaHelper.createAttributesVector(soundClassNames,
                model.getAttributes());

        Collection<Sound> sounds = Sound.query(pm, "1", detectorId);
        Instances instances = WekaHelper.getInstances(sounds, attInfo);

        // build classifier
        if (instances.numInstances() != 0)
        {
            Classifier classifier = model.getClassifier();
            classifier.buildClassifier(instances);
            model.setClassifier(classifier);
        }
        else
        {
            model.setClassifier(null);
        }
    }
}
